from langchain_agent import get_all_agent
from langgraph.graph import StateGraph,START, END
from typing_extensions import TypedDict
from langchain_core.messages.base import BaseMessage
from typing import List,Optional,Annotated
from operator import add
from langgraph.checkpoint.memory import MemorySaver

from langgraph.types import Command,interrupt
from langgraph.checkpoint.mongodb import MongoDBSaver

from mongodb_checkpoint import get_three_col
import sys

agents = get_all_agent()

class MyState(TypedDict):
    symptom: Annotated[str,add]=None
    raw_history: Annotated[str,add]=None
    structured_summary: str=None
    recipe: Annotated[Optional[str],add]=None
    diagnosis_chat_history: Annotated[Optional[str],add]=None


class SubPaserState(TypedDict):
    raw_history: Annotated[str,add]=None
    structured_summary: str=None
    symptom: Annotated[str,add]=None
    
    
class SubdiagnosisState(TypedDict):
    symptom: Annotated[str,add]=None
    structured_summary: str=None
    diagnosis_chat_history: Annotated[Optional[str],add]=None
    recipe: Annotated[Optional[str],add]=None


def history_node(state):
    msg = {"raw_history":state["raw_history"],'structured_summary':state['structured_summary']}
    history_agent = agents["history_agent"]
    res = history_agent.invoke(msg)
    last_message = res
    print('history_node last words is ',last_message)
    goto =  "parse_agent"
    return Command(
        update={
            # share internal message history of research agent with other agents
            "structured_summary":last_message,
        },
        goto=goto,
    )

def parse_node(state):
    msg = {'symptom':state['symptom'],'structured_summary':state['structured_summary']}
    parse_agent = agents['parse_agent']
    res = parse_agent.invoke(msg)
    last_message = res
    print('parse node last_message words is ',last_message)
    if "患者主诉：" in last_message:
        return Command(
            goto='sub_diagnosis_graph',
            graph=Command.PARENT,
            update={"symptom": last_message,'structured_summary':state['structured_summary']}
        )
    else:
        return Command(
            # Specify which agent to call next
            goto="parse_agent",
            # Update the graph state
            update={"symptom": last_message}
        )

def diagnosis_node(state):
    user_input = interrupt("您还有什么不舒服的吗？")
    msg = {'structured_summary':state['structured_summary'],'symptom':state['symptom'],'diagnosis_chat_history':state['diagnosis_chat_history']}
    print("msg is ",msg)
    diagnosis_agent = agents['diagnosis_agent']
    res = diagnosis_agent.invoke(msg)
    last_message = res
    print("last mes is ",last_message)
    sys.exit(0)
    if "【诊疗建议】" in last_message:
        return  Command(
            # Specify which agent to call next
            goto="recipe_agent",
            # Update the graph state
            update={"diagnosis_chat_history": last_message}
        )
    else:
        return Command(
            # Specify which agent to call next
            goto="diagnosis_agent",
            # Update the graph state
            update={"diagnosis_chat_history": last_message}
        )
    
    
def recipe_node(state):
    msg = {'diagnosis_chat_history':state['diagnosis_chat_history'],'symptom':state['symptom']}
    recipe_agent = agents['recipe_agent']
    res = recipe_agent.invoke(msg)
    last_message = res
    if "无法开药方" in last_message["result"]:
        return Command(
            # Specify which agent to call next
            goto="diagnosis_agent",
            
        )
    elif "药方" in last_message["result"]:
        return Command(
            # Specify which agent to call next
            goto="recheck_agent",
            # Update the graph state
            update={"recipe":last_message["result"]},
        )

def recheck_node(state):
    msg = {'structured_summary':state['structured_summary'],'diagnosis_chat_history':state['diagnosis_chat_history'],'recipe':state['recipe']}
    recheck_agent = agents['recheck_agent']
    res = recheck_agent.invoke(msg)
    last_message = res
    print("recheck_node is ",last_message)
    if "完全合适" in last_message:
         return Command(
            # Specify which agent to call next
            goto=END,
            # Update the graph state

        )
    else:
        return Command(
            # Specify which agent to call next
            goto="diagnosis_agent",
            # Update the graph state
            update={"symptom": last_message}
        )

def create_stateGraph():
    parent_graph_builder = StateGraph(MyState)
    sub_graph_paser_builder = StateGraph(SubPaserState)
    sub_graph_diagnosis_builder = StateGraph(SubdiagnosisState)
    
    start = START
    sub_graph_paser_builder.add_node('history_agent',history_node)
    sub_graph_paser_builder.add_node('parse_agent',parse_node)
    sub_graph_paser_builder.add_edge(start,'history_agent')
    sub_graph_paser_graph = sub_graph_paser_builder.compile()
    
    sub_graph_diagnosis_builder.add_node('diagnosis_agent',diagnosis_node)
    sub_graph_diagnosis_builder.add_node('recipe_agent',recipe_node)
    sub_graph_diagnosis_builder.add_node('recheck_agent',recheck_node)
    sub_graph_diagnosis_builder.add_edge(start,'diagnosis_agent')
    sub_graph_diagnosis_graph = sub_graph_diagnosis_builder.compile()
    
    parent_graph_builder.add_node("sub_paser_graph", sub_graph_paser_graph)
    parent_graph_builder.add_node("sub_diagnosis_graph", sub_graph_diagnosis_graph)
    parent_graph_builder.add_edge(start,"sub_paser_graph")
    
    # parent_graph = parent_graph_builder.compile()
    return parent_graph_builder

def stream_graph_updates(graph,user_input: str):
    for event in graph.stream({"messages": [{"role": "user", "content": user_input}]}):
        for value in event.values():
            print("Assistant:", value["messages"][-1].content)
 
def get_check_point():
    memory = MemorySaver()
    return memory

def test_subgraph():
    sub_graph_paser_builder = StateGraph(SubPaserState)
    
    start = START
    sub_graph_paser_builder.add_node('history_agent',history_node)
    sub_graph_paser_builder.add_node('parse_agent',parse_node)
    sub_graph_paser_builder.add_edge(start,'history_agent')
    
   
    return sub_graph_paser_builder

def sub_graph_test():
   
    sub_graph_paser_builder= test_subgraph()
    
    checkpointer = get_three_col()
    
    sub_graph_paser_graph = sub_graph_paser_builder.compile(checkpointer=checkpointer)
    config = {"configurable": {"thread_id": "1", "checkpoint_id": None}}
    res = sub_graph_paser_graph.invoke({"raw_history":long_text,"structured_summary":"病人无总结病史","symptom":"我最近有点头疼脑热"},config=config)
    # print('ssss ',type(checkpointer.get(config)["channel_values"]))
    # for message,value in checkpointer.get(config)["channel_values"].items():
    #     print(value)
    #     print(type(value))
    # print(res)
    # print(type(res))
    
if __name__=="__main__":
    # checkpointer = get_three_col()
    # checkpointer = MemorySaver()
    checkpointer = get_check_point()
    
    parent_graph_builder= create_stateGraph()
    
    parent_graph = parent_graph_builder.compile(checkpointer=checkpointer)
    
    # parent_graph = parent_graph_builder.compile()
    
    long_text =""""
    第一次就诊（2024.10.28）
医生：您好，哪里不舒服？
患者：医生，我发烧3天了，体温38.5℃，还有鼻塞、喉咙痛。
医生：有咳嗽咳痰吗？最近接触过流感病人吗？
患者：干咳没痰，同事上周也有感冒的。
医生：检查咽部充血，扁桃体不大。血常规显示病毒感染。开退烧药和抗病毒口服液，多休息多喝水。
患者：需要请假吗？
医生：建议休息3天，发热超过39℃或症状加重随时复诊。
第二次就诊（2024.11.18，间隔3周）
医生：上次用药后好些了吗？
患者：退烧了但一直流黄鼻涕，这两天又发烧到37.8℃。
医生：听诊肺部无啰音，鼻窦有压痛。可能合并细菌感染，加用抗生素，用生理盐水洗鼻。
患者：会传染给孩子吗？
医生：戴口罩避免密切接触，孩子出现症状要及时检查。
第三次就诊（2024.12.16，间隔4周）
医生：怎么又发烧了？
患者：这次是低烧37.5℃，但头痛乏力，咳嗽带黄痰。
医生：查体咽红，肺部呼吸音粗。建议拍胸片排除肺炎，近期免疫力可能较差。
患者：为什么反复感冒？
医生：需要查免疫功能和维生素D水平，平时注意锻炼和营养均衡。
    """
    message = {"messages":[{"raw_history":long_text},
        {"symptom":"我最近有点头疼"},
        {"structured_summary":"无总结病史"}]}
    
    #langgraph 的输入和agent不一样 agent需要message 而graph只需要对应的state的键就可以了
    thread_config = {"configurable": {"thread_id": 103}}
    for event in parent_graph.stream({"raw_history":long_text,"structured_summary":"病人无总结病史","symptom":"我有点不舒服","recipe":"","diagnosis_chat_history":""},config=thread_config):
        for value in event.values():
            pass
    

    exit_commands = {'quit', 'exit', 'q', '退出'}  # 可以添加更多退出命令

    while True:
        user_input = input("请输入内容: ").strip().lower()
        
        if user_input in exit_commands:
            print("感谢使用，再见！")
            break
        temp_cmd = Command(
                resume=user_input,
                update={"symptom":user_input}
            )
        for event in parent_graph.stream(temp_cmd,config=thread_config):
            for value in event.values():
                pass
        # try:
        #     temp_cmd = Command(
        #         resume=user_input,
        #         update={"symptom":user_input}
        #     )
        #     for event in parent_graph.stream(temp_cmd,config=thread_config):
        #         for value in event.values():
        #             pass
        # except Exception as e:
        #     print("error is ",e)
        #     # 处理其他输入
    